Defer the data-parallel gradient all-reduce to update() under gradient accumulation - #5099

Stacked on #5088 (base ga-bench-5060).

What

MaxTextTrainingEngine all-reduces the whole gradient tree across data replicas once per micro-batch. Only the sum over the whole optimizer step matters, so at gradient_accumulation = N it pays N times for a reduction that is correct once.

This tags the parameters reduced over the data axis right where jax.value_and_grad differentiates them. Their cotangents then come out unreduced — a per-replica partial sum that accumulates replica-locally across micro-batches — and update() reshards back to a plain spec, which is what emits the single cross-replica all-reduce.

It is the same trick gradient_accumulation.py already plays for the pre-train path, applied across the engine's separate jax.jit dispatches rather than inside one lax.scan.

Results — qwen3-0.6b on 4× v6e

data=4, fsdp=1, micro-batch 8×1024, shard_mode=explicit, SGD, no clipping, remat=none. Median steady-state step time over 20 post-warmup steps, untraced:

Where it went, at GA=8: fwd_bwd drops 59.6 → 37.4 ms per micro-batch, and update does not pay it back on the host clock (82.9 → 60.8 ms). These two are host-side dispatch medians — the engine's calls are asynchronous, so they carry queue backpressure and do not sum to the step time; the end-to-end median in the table is the number to trust. They locate the win, they don't account for it.

The optimized HLO for the real model is unambiguous:

Losses agree to within 4.5e-5 relative across all 23 steps of every A/B pair (max over GA ∈ {1, 4, 8}) — float32 reassociation of the same sum, since the cross-replica addition moves from before the micro-batch sum to after it. On a 4-device CPU mesh, where the reassociation is exact, they are bit-identical.

When it engages

_deferred_all_reduce_shardings returns (None, None) — the untagged status quo — unless all of:

1. shard_mode == EXPLICIT;
2. the mesh's axis types are all Explicit (a caller can hand the engine a bare jax.sharding.Mesh(...) regardless of shard_mode, and the tags are rejected on Auto axes);
3. data is the only mesh axis of size > 1 that activation_batch resolves to.

(3) is not conservatism. A gradient contracts over the batch and JAX requires the unreduced set to be exactly the contracted axes, so with fsdp on the batch too it rejects the backward pass outright:

ShardingTypeError: out_sharding's unreduced axes should be equal to the contracting specs.
Got unreduced axes=frozenset({'data'}) and contracting spec=(('data', 'fsdp'), None)

and widening the tag to fsdp is not available either, since the parameters are sharded over it. Verified on a data=2 × fsdp=2 mesh.

What else had to change

• layers/normalizations.py — _align_scale_with_normalized_axis indexed spec[...], which a tagged spec refuses (ValueError: Using pspec[...] is dangerous when PartitionSpec has non-empty unreduced/reduced set). It reads spec.partitions now and carries the tags onto the new spec. Every run crashed here before this; it is the same fix the pre-train path needed.
• utils/sharding.py — new batch_mesh_axes(mesh, rules) for condition (3).
• Accumulator lifetime — an accumulator can outlive the shardings it was produced under. A checkpoint holds the reduced total (Orbax cannot serialize an unreduced array at all: device_indices_map is undefined for one), and a recompile can flip the deferral on or off. So save_checkpoint reduces on the way out, and both restore_checkpoint and _compile_for_batch move a live accumulator back onto whatever the kernels now expect. Both directions are exact: resharding away from unreduced runs the pending all-reduce, and device_put onto it leaves the value on one replica and zeros the rest, so the deferred all-reduce reproduces it.

_accumulated_denominator is deliberately left plain — float() on an unreduced scalar raises, and its per-micro-batch all-reduce is a single f32[].

Second commit: the gate missed tensor parallelism

Found while benchmarking #5104 against a pure-TP mesh. The gate refused a second mesh axis on the batch dimension, which catches fsdp. tensor gets to the same contradiction through the feature dimension of the same activation, which a batch-axis check cannot see, so qwen3-0.6b at dp2 × tp2 on 4× v6e died on the first micro-batch:

ShardingTypeError: out_sharding's unreduced axes should be equal to the contracting specs.
Got unreduced axes=frozenset({'data'}) and contracting spec=('data', None, 'tensor')

expert, context and tensor_sequence shard contracted dimensions too and would have gone the same way. Enumerating the axes known to break is how tensor was missed, so the rule is now the blunt one — data alone above size 1, or no deferral. Nothing that worked before is lost, since none of those meshes ran.

Two tests come with it: the gate declines at dp2 × tp2 (fails without the fix), and a tensor-parallel mesh completes a step and moves its weights with the tag off. The second does not reproduce the crash — the toy model shards plenty over tensor and still traces clean on CPU; what raised the error was qwen3-0.6b's attention kernels on TPU, verified there before and after.

Tests

tests/post_training/unit/maxtext_engine_deferred_all_reduce_test.py, 13 cases on a 4-device CPU mesh with a real (tiny) MaxText decoder:

• the gate opens on a purely data-parallel explicit mesh and declines under auto shard mode, on an Auto-axis mesh, with no data replicas, with fsdp on the batch, with tensor on the features, and for a parameter already sharded over data;
• a tensor-parallel mesh, with the gate declining, still completes a step and moves its weights;
• the micro-batch kernels' optimized HLO carries no array all-reduce, with a vacuity guard that update()'s does;
• the mirror-image assertion with the deferral withheld, proving the probe can fail;
• weights match a non-deferred run;
• an unreduced accumulator survives a mid-step checkpoint round trip and finishes the step to the same weights as an uninterrupted run. (This one found a real bug: restoring does not recompile, so the restored plain total reached kernels expecting an unreduced one.)

Existing maxtext_engine_test.py + maxtext_engine_constructor_test.py: 54 passed.

Follow-ups, not in scope here

• ZeRO-1 is a silent no-op in the engine. shard_optimizer_over_data is read only by gradient_accumulation.py, which the engine does not go through. It is orthogonal to this change — the win above needs no ZeRO-1 — but wiring it into the engine is a separate and much larger piece of work.
• The engine's explicit-sharding support still rests on tunix_adapter.py:67's process-wide with_sharding_constraint → reshard monkeypatch.
• Reproducing the benchmark needs three flags added to the perf_parity rig from Match Tunix peft_trainer_v2 performance in MaxTextTrainingEngine #5060 (--dp, --shard-mode, --no-defer) plus building the mesh with maxtext_utils.get_mesh_from_config instead of a bare jax.sharding.Mesh, which is what actually sets AxisType.Explicit. Those files are not on this branch, so the change is not included here; the command was python qwen3_engine_profile.py --ga 8 --dp 4 --fsdp 1 --shard-mode explicit --no-trace [--no-defer].